import { requestBailian } from './bailian.js';
import { executeSql } from './mysql.js';
import { chatSSE } from './qwen.js';
import { getContentBetweenStart, autoFormatValue } from '../tools/str.js';
import nodeSqlParse from 'node-sql-parser';
import { logger } from '../tools/index.js';
import { render } from '../prompts/index.js';
const { Parser } = nodeSqlParse;

function setSelectFields(astItem, tableSchema) {
    const selectFields = tableSchema.columns;
    // 构造新的 SELECT 部分
    astItem.columns = selectFields.map(field => ({
        expr: {
            type: 'column_ref',
            table: null,
            column: field.col_name
        },
        as: null
    }));
}

function appendContextConditions(astItem, context) {
    const bookId = context.book.id;
    const createUserId = context.user.id;
    // const additionalConditions = {
    //     type: 'binary_expr',
    //     operator: 'AND',
    //     left: {
    //         type: 'binary_expr',
    //         operator: '=',
    //         left: { type: 'column_ref', table: null, column: 'book_id' },
    //         right: { type: 'string', value: bookId },
    //     },
    //     right: {
    //         type: 'binary_expr',
    //         operator: '=',
    //         left: { type: 'column_ref', table: null, column: 'create_account_no' },
    //         right: { type: 'string', value: createUserId },
    //     }
    // };
    const additionalConditions = {
        type: 'binary_expr',
        operator: '=',
        left: { type: 'column_ref', table: null, column: 'book_id' },
        right: { type: 'string', value: bookId },
    };

    // 如果原始 SQL 中已经有 WHERE 条件，则需要构建一个逻辑 AND 表达式
    if (astItem.where) {
        astItem.where = {
            type: 'binary_expr',
            operator: 'AND',
            left: astItem.where, // 原有的 WHERE 条件
            right: additionalConditions, // 新的额外条件
        };
    } else {
        // 如果原始 SQL 中没有 WHERE 条件，则可以直接添加
        astItem.where = additionalConditions;
    }
}

// 自动格式化SQL适应查询明细
const sqlProcess = (sql, tableSchema, context) => {
    try {
        const parser = new Parser();
        let ast = parser.astify(sql);
        ast = Array.isArray(ast) ? ast : [ast];
        if (ast && ast.length === 1) {
            const astItem = ast[0];
            if (astItem.type === 'select') {
                // 设置查询字段
                setSelectFields(astItem, tableSchema);
                // 额外添加的条件
                appendContextConditions(astItem, context);
                // 设置其他条件
                astItem.groupby = null;
                astItem.orderby = [
                    {
                        expr: {
                            type: 'column_ref',
                            table: null, // 如果有明确的表名，可以替换null为表名
                            column: 'id'
                        },
                        type: 'DESC'
                    }
                ];
            }
            return parser.sqlify(ast);
        }
        return sql;
    } catch (err) {
        console.error(`sqlProcess error, sql = ${sql}, error = `, err);
        return sql;
    }
}

async function n2sqlByBailian(prompt, tableSchema, context) {
    const result = await requestBailian(process.env.BAILIAN_N2SQL_APP_ID, prompt, {
        "sqlInput": {
            "synonym_infos": "",
            "schema_infos": [tableSchema]
        }
    }, context);
    return result.output.text;
}

async function n2sqlByLocalPrompt(prompt, tableSchema, context, examples = []) {
    const finalPrompt = render('n2data.handlebars', {
        schema: {
            table_id: tableSchema.table_id,
            table_desc: tableSchema.table_desc,
            columns: tableSchema.columns.filter(col => col.useFilter)
        },
        question: prompt,
        examples: examples,
    });
    const finalAnswer = await chatSSE(finalPrompt, () => { }, context);
    return finalAnswer.output.text.replace('Output:', '').trim();
}

export default async function (prompt, tableSchema, context, maxSize, examples) {
    // N2SQL
    const tableColumns = tableSchema.columns;
    // const sql = await n2sqlByBailian(prompt, tableSchema, context);
    const sql = await n2sqlByLocalPrompt(prompt, tableSchema, context, examples);
    const preProcessSql = sqlProcess(sql, tableSchema, context);
    logger.info(`n2data start, prompt = ${prompt}, originalSql = ${sql}, preProcessSql = ${preProcessSql}`);

    // 执行统计SQL
    const finalCountSql = `SELECT COUNT(*) AS cnt FROM (${preProcessSql}) t`;
    const countData = await executeSql(finalCountSql);
    if (countData.length > 0 && countData[0].cnt > maxSize) {
        throw new Error(`最多支持${maxSize}行数据查出，待分析的数据过多，小钱钱无法操作。`);
    }

    // 执行SQL获取数据
    const finalSql = `SELECT * FROM (${preProcessSql}) t LIMIT ${maxSize}`;
    const resultData = await executeSql(finalSql);
    const formatResultData = resultData.map(row => {
        const processRow = {};
        tableColumns.forEach(column => {
            const columnDesc = getContentBetweenStart(column.col_caption, '', '，') || column.col_caption;
            processRow[columnDesc] = autoFormatValue(row[column.col_name]);
        });
        return processRow;
    });
    return [formatResultData, sql]
}
